"""
知识库模块
提供文档管理和检索功能，支持按角色查询和文件上传
"""

from typing import List, Dict, Any, Optional
import os
import time
import logging
import shutil
import re

# 设置日志
logger = logging.getLogger("KnowledgeBase")
logger.setLevel(logging.INFO)

class KnowledgeBase:
    def __init__(self, knowledge_dir="knowledge"):
        """初始化知识库
        
        Args:
            knowledge_dir: 知识库目录
        """
        self.logger = logging.getLogger("KnowledgeBase")
        self.logger.setLevel(logging.INFO)
        
        # 知识库目录
        self.knowledge_dir = knowledge_dir
        if not os.path.exists(self.knowledge_dir):
            os.makedirs(self.knowledge_dir)
            self.logger.info(f"创建知识库目录: {self.knowledge_dir}")
        
        # 文档列表，按角色分类
        self.documents = {}
        self.all_documents = []
        
        # 加载文档
        self.load_documents()
    
    def load_documents(self):
        """加载知识库文件"""
        self.documents = {}
        self.all_documents = []
        
        # 遍历知识库目录
        for root, dirs, files in os.walk(self.knowledge_dir):
            # 获取当前目录相对于知识库目录的路径
            rel_path = os.path.relpath(root, self.knowledge_dir)
            role_id = rel_path if rel_path != "." else "general"
            
            # 确保角色文档列表存在
            if role_id not in self.documents:
                self.documents[role_id] = []
            
            for file in files:
                if file.endswith('.txt') or file.endswith('.md') or file.endswith('.pdf'):
                    file_path = os.path.join(root, file)
                    try:
                        with open(file_path, 'r', encoding='utf-8') as f:
                            content = f.read()
                            
                            # 分段处理
                            paragraphs = content.split('\n\n')
                            for para in paragraphs:
                                if para.strip():
                                    doc = {
                                        'content': para.strip(),
                                        'source': file_path,
                                        'role_id': role_id
                                    }
                                    self.documents[role_id].append(doc)
                                    self.all_documents.append(doc)
                    except Exception as e:
                        self.logger.error(f"加载文件失败: {file_path}, 错误: {str(e)}")
        
        self.logger.info(f"知识库加载完成，共 {len(self.all_documents)} 个文档段落，{len(self.documents)} 个角色")
    
    def add_document(self, content: str, role_id: str = "general", source: str = "用户上传"):
        """添加文档到知识库
        
        Args:
            content: 文档内容
            role_id: 角色ID，如果不存在则创建
            source: 文档来源
            
        Returns:
            是否添加成功
        """
        # 确保角色目录存在
        role_dir = self.knowledge_dir if role_id == "general" else os.path.join(self.knowledge_dir, role_id)
        if not os.path.exists(role_dir):
            os.makedirs(role_dir)
            self.logger.info(f"创建角色知识库目录: {role_dir}")
        
        # 保存到文件
        file_path = os.path.join(role_dir, f"{int(time.time())}.txt")
        try:
            with open(file_path, 'w', encoding='utf-8') as f:
                f.write(content)
            
            # 确保角色文档列表存在
            if role_id not in self.documents:
                self.documents[role_id] = []
            
            # 分段处理
            paragraphs = content.split('\n\n')
            for para in paragraphs:
                if para.strip():
                    doc = {
                        'content': para.strip(),
                        'source': source,
                        'role_id': role_id
                    }
                    self.documents[role_id].append(doc)
                    self.all_documents.append(doc)
            
            self.logger.info(f"文档添加成功，角色: {role_id}，共 {len(paragraphs)} 个段落")
            return True
        except Exception as e:
            self.logger.error(f"添加文档失败: {str(e)}")
            return False
    
    def upload_file(self, file_path: str, role_id: str = "general", file_name: Optional[str] = None) -> bool:
        """上传文件到知识库
        
        Args:
            file_path: 文件路径
            role_id: 角色ID，如果不存在则创建
            file_name: 文件名，如果不提供则使用原文件名
            
        Returns:
            是否上传成功
        """
        try:
            # 确保角色目录存在
            role_dir = self.knowledge_dir if role_id == "general" else os.path.join(self.knowledge_dir, role_id)
            if not os.path.exists(role_dir):
                os.makedirs(role_dir)
                self.logger.info(f"创建角色知识库目录: {role_dir}")
            
            # 确定目标文件名
            if not file_name:
                file_name = os.path.basename(file_path)
            
            # 目标文件路径
            target_path = os.path.join(role_dir, file_name)
            
            # 复制文件
            shutil.copy2(file_path, target_path)
            
            # 如果是文本文件，加载内容到知识库
            if file_name.endswith('.txt') or file_name.endswith('.md'):
                try:
                    with open(target_path, 'r', encoding='utf-8') as f:
                        content = f.read()
                        
                        # 确保角色文档列表存在
                        if role_id not in self.documents:
                            self.documents[role_id] = []
                        
                        # 分段处理
                        paragraphs = content.split('\n\n')
                        for para in paragraphs:
                            if para.strip():
                                doc = {
                                    'content': para.strip(),
                                    'source': target_path,
                                    'role_id': role_id
                                }
                                self.documents[role_id].append(doc)
                                self.all_documents.append(doc)
                except Exception as e:
                    self.logger.error(f"加载文件内容失败: {target_path}, 错误: {str(e)}")
            
            self.logger.info(f"文件上传成功: {target_path}")
            return True
        except Exception as e:
            self.logger.error(f"文件上传失败: {str(e)}")
            return False
    
    def search(self, query: str, role_id: Optional[str] = None, top_k: int = 5) -> List[Dict[str, Any]]:
        """搜索知识库
        
        Args:
            query: 查询文本
            role_id: 角色ID，如果提供则只搜索该角色的知识库
            top_k: 返回结果数量
            
        Returns:
            匹配的文档列表
        """
        # 确定搜索范围
        if role_id and role_id in self.documents:
            search_docs = self.documents[role_id]
        else:
            search_docs = self.all_documents
        
        if not search_docs:
            self.logger.warning(f"知识库为空，无法搜索，角色ID: {role_id}")
            return []
        
        # 使用简单文本匹配
        query_terms = re.findall(r'\w+', query.lower())
        
        if not query_terms:
            self.logger.warning("查询词为空，返回前几个文档")
            return search_docs[:top_k]
        
        results = []
        for doc in search_docs:
            content = doc['content'].lower()
            score = sum(1 for term in query_terms if term in content)
            # 确保至少返回一些结果
            results.append({
                'content': doc['content'],
                'source': doc.get('source', 'unknown'),
                'role_id': doc.get('role_id', 'general'),
                'score': score
            })
        
        # 按得分排序
        results.sort(key=lambda x: x['score'], reverse=True)
        
        self.logger.info(f"搜索完成，角色ID: {role_id}，找到 {len(results)} 个结果")
        return results[:top_k]